from typing import Dict, Any, List

from langchain_community.chat_models import ChatTongyi
from langchain_core.callbacks import BaseCallbackHandler
import os

from langchain_core.outputs import LLMResult

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"


class MyCustomSyncHandler(BaseCallbackHandler):
    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None:
        """Called when language model starts running."""
        print("Language model started running.")

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Called when language model ends running."""
        print("Language model finished running.")

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Called when a new token is generated by the language model."""
        print(token, end='', flush=True)


if __name__ == "__main__":

    llm = ChatTongyi(
        streaming=True,
        callbacks=[MyCustomSyncHandler()],
    )
    llm.invoke("请详细描述java和python的区别")
